from hashlib import md5
from os import urandom
from base64 import urlsafe_b64encode
from flask import Blueprint, request
from time import time
from config import Config
from utils import respond, Cache, ERRORCODES
from jwt import encode as jwt_encode
from jwt import decode as jwt_decode


bp = Blueprint("AuthComponent", __name__, url_prefix="/auth/")


def generate_token(byte_length=16):
    random_bytes = urandom(byte_length)
    token = urlsafe_b64encode(random_bytes).rstrip(b"=").decode("utf-8")
    return token


@bp.route("/get_ticket", methods=["GET"])
def get_ticket():
    device_id = request.args.get("device_id")
    if not device_id:
        return respond(None, "缺少设备ID。", ERRORCODES.MISSING_PARAMETERS)
    elif device_id not in Config.DEVICES:
        return respond(None, "未知的设备ID。", ERRORCODES.NOT_FOUND)

    ticket = generate_token(24)
    Cache.set(ticket, device_id, Config.TICKET_EXPIRES)
    return respond({"ticket": ticket})


@bp.route("/get_token", methods=["POST"])
def get_token():
    device_id, ticket, signature = (
        request.json.get("device_id"),
        request.json.get("ticket"),
        request.json.get("signature"),
    )
    if not device_id or not ticket or not signature:
        return respond(None, "缺少参数。", ERRORCODES.BAD_REQUEST)
    elif device_id not in Config.DEVICES or Cache.get(ticket) != device_id:
        return respond(
            None, "无效/过期的device_id或ticket。", ERRORCODES.BAD_REQUEST
        )
    elif (
        signature
        != md5(
            (ticket + device_id + Config.DEVICES[device_id]).encode("utf-8")
        ).hexdigest()
    ):
        return respond(None, "无效的签名。", ERRORCODES.UNMATCHED_SIGNATURE)

    payload = {
        "refresh_remaining": Config.JWT_REFRESH_TIMES,
        "device_id": device_id,
        "exp": int(time()) + Config.JWT_EXPIRES,
    }
    return respond(
        {"token": jwt_encode(payload, Config.JWT_SECRET_KEY, algorithm="HS256")}
    )


@bp.route("/refresh_token", methods=["POST"])
def refresh_token():
    device_id, token, signature = (
        request.json.get("device_id"),
        request.json.get("token"),
        request.json.get("signature"),
    )

    if not device_id or not token or not signature:
        return respond(None, "缺少参数。", ERRORCODES.BAD_REQUEST)
    elif (
        signature
        != md5(
            (token + device_id + Config.DEVICES[device_id]).encode("utf-8")
        ).hexdigest()
    ):
        return respond(None, "无效的签名。", ERRORCODES.UNMATCHED_SIGNATURE)

    try:
        old_token = jwt_decode(token, Config.JWT_SECRET_KEY, algorithms=["HS256"])
    except:  # Exception as e: 
        # raise e
        return respond(None, "无效的token。", ERRORCODES.UNAUTHORIZED)

    if old_token["refresh_remaining"] <= 0:
        return respond(None, "刷新次数已用尽。", ERRORCODES.EXHAUSTED_REFRESH_TIMES)

    payload = {
        "refresh_remaining": old_token["refresh_remaining"] - 1,
        "device_id": device_id,
        "exp": int(time()) + Config.JWT_EXPIRES,
    }

    return respond(
        {"token": jwt_encode(payload, Config.JWT_SECRET_KEY, algorithm="HS256")}
    )